from openni import openni2
import numpy as np
import cv2
import time

# ======================= 参数设置 ==========================
WINDOW_NAME_DEPTH = 'Depth Image'
WINDOW_NAME_COLOR = 'Color Image'
COLOR_MAP_TYPE = 8      # 0~11 不同色彩映射，8色彩鲜艳
ALPHA_VALUE = 0.17
MAX_DISTANCE_CM = 800   # 最大有效距离（cm）
FONT = cv2.FONT_HERSHEY_SIMPLEX
FONT_SCALE = 0.5
FONT_COLOR = (0, 0, 0)
FONT_THICKNESS = 1
last_click_time = 0
click_x, click_y = -1, -1
distance_text = ""
# ==========================================================

# 鼠标点击获取深度
def mousecallback(event, x, y, flags, param):
    global click_x, click_y, distance_text, last_click_time
    if event == cv2.EVENT_LBUTTONDOWN:
        click_x, click_y = x, y
        distance = dpt[y, x] / 10.0  # mm 转 cm
        distance_text = f"Dis: {distance:.2f} cm"
        last_click_time = time.time()

if __name__ == "__main__":
    try:
        # 初始化 OpenNI 深度设备
        openni2.initialize()
        dev = openni2.Device.open_any()
        print(dev.get_device_info())

        depth_stream = dev.create_depth_stream()
        dev.set_image_registration_mode(True)
        depth_stream.start()

        # 初始化 RGB 摄像头（cv2）
        cap = cv2.VideoCapture(0)
        if not cap.isOpened():
            print("RGB camera not found")
            exit()

        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        cap.set(cv2.CAP_PROP_FPS, 30)

        # 创建窗口
        cv2.namedWindow(WINDOW_NAME_DEPTH)
        cv2.namedWindow(WINDOW_NAME_COLOR)
        cv2.setMouseCallback(WINDOW_NAME_DEPTH, mousecallback)

        print("Press SPACE to quit.")

        # 实时循环
        while True:
            # ---------------- 深度流 ----------------
            frame = depth_stream.read_frame()
            dframe_data = np.array(frame.get_buffer_as_triplet()).reshape([480, 640, 2])
            dpt1 = np.asarray(dframe_data[:, :, 0], dtype='float32')
            dpt2 = np.asarray(dframe_data[:, :, 1], dtype='float32')
            dpt2 *= 255
            dpt = dpt1 + dpt2

            dim_gray = cv2.convertScaleAbs(dpt, alpha=ALPHA_VALUE)
            depth_colormap = cv2.applyColorMap(dim_gray, COLOR_MAP_TYPE)

            # 点击显示距离
            if click_x >= 0 and click_y >= 0 and (time.time() - last_click_time) < 5:
                depth_colormap = cv2.putText(
                    depth_colormap, distance_text, (click_x, click_y),
                    FONT, FONT_SCALE, FONT_COLOR, FONT_THICKNESS, cv2.LINE_AA
                )

            cv2.imshow(WINDOW_NAME_DEPTH, depth_colormap)

            # ---------------- RGB 流 ----------------
            ret, frame_rgb = cap.read()
            if ret:
                # 镜像还原
                frame_rgb = cv2.flip(frame_rgb, 1)
                cv2.imshow(WINDOW_NAME_COLOR, frame_rgb)

            # ---------------- 退出控制 ----------------
            key = cv2.waitKey(1)
            if key & 0xFF == ord(' '):  # 空格键退出
                break

        # 释放资源
        depth_stream.stop()
        dev.close()
        openni2.unload()
        cap.release()
        cv2.destroyAllWindows()

    except Exception as e:
        print(f"An error occurred: {e}")
